﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using Jarvis.Learners;

namespace Jarvis.Learners.DecisionTree
{
    class DecisionTreeLearner : Learner
    {
        protected TreeNode root;

        public DecisionTreeLearner()
        {

        }

        protected virtual String GetFeatureForExamples(List<Example> node_examples)
        {
            foreach (Example e in node_examples)
            {
                foreach (String key in e.state.Keys)
                    return key;
            }

            return null;
        }

        protected TreeNode TreeNodeForExamples(List<Example> node_examples)
        {
            if (node_examples == null || node_examples.Count == 0)
                return null;

            node_examples = CopyExamplesList(node_examples);

            Boolean all_same = true;
            String label = null;

            foreach (Example e in node_examples)
            {
                if (!e.label.Equals(label))
                {
                    if (label != null)
                        all_same = false;

                    label = e.label;
                }
            }

            TreeNode node = new TreeNode();

            if (all_same)
            {
                node.label = label;
                node.examples_count = (Double) node_examples.Count;
            }
            else
            {
                node.examples_count = (Double)node_examples.Count;
                String feature = this.GetFeatureForExamples(node_examples);

                if (feature == null)
                    return null;

                node.feature_name = feature;

                Dictionary<Object, List<Example>> example_lists = new Dictionary<object, List<Example>>();
                List<Example> unknowns = new List<Example>();

                if (node_examples[0].splits.ContainsKey(feature))
                {
                    double split = node_examples[0].splits[feature];

                    node.split = split;

                    foreach (Example e in node_examples)
                    {
                        if (e.state.ContainsKey(feature))
                        {
                            Double value = (Double)e.state[feature];

                            String new_value = "<";

                            if (value >= split)
                                new_value = ">=";

                            List<Example> matches = null;

                            if (example_lists.ContainsKey(new_value))
                                matches = example_lists[new_value];
                            else
                            {
                                matches = new List<Example>();
                                example_lists[new_value] = matches;
                            }

                            e.state.Remove(feature);

                            matches.Add(e);
                        }
                        else
                            unknowns.Add(e);

                        foreach (Object key in example_lists.Keys)
                        {
                            TreeNode child = this.TreeNodeForExamples(example_lists[key]);

                            if (child != null)
                            {
                                if (example_lists.ContainsKey(key))
                                    child.examples_count = (Double)example_lists[key].Count;
                                else
                                    child.examples_count = 1;

                                node.children[key] = child;
                            }
                        }
                    }
                }
                else
                {
                    if (node_examples[0].subitems.ContainsKey(feature))
                    {
                        object subitem = node_examples[0].subitems[feature];

                        node.subitem = subitem;

                        foreach (Example e in node_examples)
                        {
                            if (e.state.ContainsKey(feature))
                            {
                                List<object> value = (List<object>)e.state[feature];

                                String new_value = "Does not contain";

                                if (value.Contains(subitem))
                                    new_value = "Contains";

                                List<Example> matches = null;

                                if (example_lists.ContainsKey(new_value))
                                    matches = example_lists[new_value];
                                else
                                {
                                    matches = new List<Example>();
                                    example_lists[new_value] = matches;
                                }

                                e.state.Remove(feature);

                                matches.Add(e);
                            }
                            else
                                unknowns.Add(e);

                            foreach (Object key in example_lists.Keys)
                            {
                                TreeNode child = this.TreeNodeForExamples(example_lists[key]);

                                if (child != null)
                                {
                                    child.examples_count = (Double)example_lists[key].Count;
                                    node.children[key] = child;
                                }
                            }
                        }
                    }
                    else
                    {
                        foreach (Example e in node_examples)
                        {
                            if (e.state.ContainsKey(feature))
                            {
                                Object value = e.state[feature];

                                List<Example> matches = null;


                                if (example_lists.ContainsKey(value))
                                    matches = example_lists[value];
                                else
                                {
                                    matches = new List<Example>();
                                    example_lists[value] = matches;
                                }

                                matches.Add(e);
                            }
                            else
                                unknowns.Add(e);

                            e.state.Remove(feature);
                        }

                        foreach (Object key in example_lists.Keys)
                        {
                            if (key != null)
                            {
                                TreeNode child = this.TreeNodeForExamples(example_lists[key]);

                                if (child != null)
                                {
                                    child.examples_count = (Double)example_lists[key].Count;
                                    node.children[key] = child;
                                }
                            }
                        }
                    }
                }

                if (unknowns.Count > 0)
                {
                    node.unknown_value = this.TreeNodeForExamples(unknowns); 
                }
            }

            return node;
        }

        public override void Train()
        {
            Console.WriteLine("decision tree train");

            Console.WriteLine("---- model '" + this.name + "' (" + examples.Count + " examples) ----");
            root = this.TreeNodeForExamples(examples);

            Console.WriteLine(root);
            Console.WriteLine("-- end model --");
        }

        public override String ToString()
        {
            try
            {
                return "Decision Tree Learner: " + this.name + Environment.NewLine + root.ToString()
                    + Environment.NewLine + examples.Count + " training examples.";
            }
            catch (NullReferenceException)
            {
                return "Decision Tree Learner: " + this.name + Environment.NewLine + "0 training examples.";
            }
        }

        public override String GetPrediction(Dictionary<String, object> state)
        {
            if (root != null)
                return root.GetLabel(state);
            else
                return "Unknown";
        }
    }
}
